# -*- coding: utf-8 -*-
"""
检索方法基类
定义所有检索方法的通用接口
"""

from abc import ABC, abstractmethod
from typing import Dict, List, Any, Optional, Tuple
import time
from loguru import logger


class RetrievalResult:
    """检索结果类"""
    
    def __init__(self, 
                 question: str,
                 answer: str,
                 confidence: float = 0.0,
                 response_time: float = 0.0,
                 retrieved_docs: List[Dict] = None,
                 metadata: Dict = None):
        """
        初始化检索结果
        
        Args:
            question: 原始问题
            answer: 生成的答案
            confidence: 置信度分数 (0-1)
            response_time: 响应时间（秒）
            retrieved_docs: 检索到的文档列表
            metadata: 额外的元数据
        """
        self.question = question
        self.answer = answer
        self.confidence = confidence
        self.response_time = response_time
        self.retrieved_docs = retrieved_docs or []
        self.metadata = metadata or {}
    
    def to_dict(self) -> Dict[str, Any]:
        """转换为字典格式"""
        return {
            'question': self.question,
            'answer': self.answer,
            'confidence': self.confidence,
            'response_time': self.response_time,
            'retrieved_docs': self.retrieved_docs,
            'metadata': self.metadata
        }


class BaseRetriever(ABC):
    """检索方法基类"""
    
    def __init__(self, name: str, config: Dict[str, Any] = None):
        """
        初始化检索器
        
        Args:
            name: 检索器名称
            config: 配置参数
        """
        self.name = name
        self.config = config or {}
        self.is_initialized = False
        
        # 性能统计
        self.total_queries = 0
        self.total_time = 0.0
        self.success_count = 0
        
        logger.info(f"🔧 初始化检索器: {self.name}")
    
    @abstractmethod
    def initialize(self, data: List[Dict[str, Any]]) -> bool:
        """
        初始化检索器（建立索引等）
        
        Args:
            data: 训练/索引数据
        
        Returns:
            bool: 初始化是否成功
        """
        pass
    
    @abstractmethod
    def retrieve(self, question: str, top_k: int = 5) -> RetrievalResult:
        """
        执行检索
        
        Args:
            question: 查询问题
            top_k: 返回结果数量
        
        Returns:
            RetrievalResult: 检索结果
        """
        pass
    
    def batch_retrieve(self, questions: List[str], top_k: int = 5) -> List[RetrievalResult]:
        """
        批量检索
        
        Args:
            questions: 问题列表
            top_k: 返回结果数量
        
        Returns:
            List[RetrievalResult]: 检索结果列表
        """
        results = []
        for question in questions:
            result = self.retrieve(question, top_k)
            results.append(result)
        return results
    
    def _measure_time(self, func, *args, **kwargs) -> Tuple[Any, float]:
        """
        测量函数执行时间
        
        Args:
            func: 要执行的函数
            *args: 函数参数
            **kwargs: 函数关键字参数
        
        Returns:
            Tuple: (函数返回值, 执行时间)
        """
        start_time = time.time()
        result = func(*args, **kwargs)
        end_time = time.time()
        execution_time = end_time - start_time
        
        # 更新统计信息
        self.total_queries += 1
        self.total_time += execution_time
        
        return result, execution_time
    
    def get_performance_stats(self) -> Dict[str, Any]:
        """
        获取性能统计信息
        
        Returns:
            Dict: 性能统计数据
        """
        avg_time = self.total_time / self.total_queries if self.total_queries > 0 else 0
        success_rate = self.success_count / self.total_queries if self.total_queries > 0 else 0
        
        return {
            'name': self.name,
            'total_queries': self.total_queries,
            'total_time': self.total_time,
            'average_time': avg_time,
            'success_count': self.success_count,
            'success_rate': success_rate
        }
    
    def reset_stats(self) -> None:
        """
        重置性能统计
        """
        self.total_queries = 0
        self.total_time = 0.0
        self.success_count = 0
    
    def calculate_precision(self, predicted: str, actual: str) -> float:
        """
        计算精确度（简单实现）
        
        Args:
            predicted: 预测答案
            actual: 实际答案
        
        Returns:
            float: 精确度分数 (0-1)
        """
        if not predicted or not actual:
            return 0.0
        
        # 简单的字符串匹配
        predicted_words = set(predicted.lower().split())
        actual_words = set(actual.lower().split())
        
        if not actual_words:
            return 0.0
        
        intersection = predicted_words.intersection(actual_words)
        return len(intersection) / len(actual_words)
    
    def calculate_recall(self, predicted: str, actual: str) -> float:
        """
        计算召回率（简单实现）
        
        Args:
            predicted: 预测答案
            actual: 实际答案
        
        Returns:
            float: 召回率分数 (0-1)
        """
        if not predicted or not actual:
            return 0.0
        
        # 简单的字符串匹配
        predicted_words = set(predicted.lower().split())
        actual_words = set(actual.lower().split())
        
        if not predicted_words:
            return 0.0
        
        intersection = predicted_words.intersection(actual_words)
        return len(intersection) / len(predicted_words)
    
    def calculate_f1_score(self, predicted: str, actual: str) -> float:
        """
        计算F1分数
        
        Args:
            predicted: 预测答案
            actual: 实际答案
        
        Returns:
            float: F1分数 (0-1)
        """
        precision = self.calculate_precision(predicted, actual)
        recall = self.calculate_recall(predicted, actual)
        
        if precision + recall == 0:
            return 0.0
        
        return 2 * (precision * recall) / (precision + recall)
    
    def evaluate_result(self, result: RetrievalResult, ground_truth: str) -> Dict[str, float]:
        """
        评估检索结果
        
        Args:
            result: 检索结果
            ground_truth: 标准答案
        
        Returns:
            Dict: 评估指标
        """
        precision = self.calculate_precision(result.answer, ground_truth)
        recall = self.calculate_recall(result.answer, ground_truth)
        f1_score = self.calculate_f1_score(result.answer, ground_truth)
        
        return {
            'precision': precision,
            'recall': recall,
            'f1_score': f1_score,
            'response_time': result.response_time,
            'confidence': result.confidence
        }
    
    def __str__(self) -> str:
        return f"{self.__class__.__name__}(name='{self.name}', initialized={self.is_initialized})"
    
    def __repr__(self) -> str:
        return self.__str__()